#!/usr/bin/env python3
"""
综合建筑安全检测系统
支持多种安全违规行为检测：安全帽、工服、抽烟、接打电话、烟火等
"""

import os
import cv2
import numpy as np
from ultralytics import YOLO
import matplotlib.pyplot as plt
import argparse
from pathlib import Path
import json
from datetime import datetime

class ComprehensiveSafetyDetector:
    def __init__(self, model_path="runs/detect/comprehensive_safety/weights/best.pt"):
        """
        初始化综合安全检测器
        
        Args:
            model_path: 训练好的模型路径
        """
        self.model_path = model_path
        self.model = None
        
        # 扩展的类别定义 - 支持多种安全检测
        self.class_names = [
            'Person',           # 0 - 人员
            'Helmet',           # 1 - 佩戴安全帽
            'No-Helmet',        # 2 - 未佩戴安全帽
            'Safety-Vest',      # 3 - 佩戴安全背心/工服
            'No-Safety-Vest',   # 4 - 未佩戴安全背心/工服
            'Smoking',          # 5 - 抽烟行为
            'Phone-Call',       # 6 - 接打电话
            'Fire',             # 7 - 火焰
            'Smoke',            # 8 - 烟雾
            'Danger-Zone',      # 9 - 危险区域
            'Safety-Equipment', # 10 - 安全设备
            'Vehicle'           # 11 - 工程车辆
        ]
        
        # 颜色配置 - 不同风险级别用不同颜色
        self.colors = {
            'Person': (255, 255, 0),           # 青色 - 人员
            'Helmet': (0, 255, 0),             # 绿色 - 安全（佩戴安全帽）
            'No-Helmet': (0, 0, 255),          # 红色 - 危险（未佩戴安全帽）
            'Safety-Vest': (0, 255, 0),        # 绿色 - 安全（佩戴工服）
            'No-Safety-Vest': (0, 165, 255),   # 橙色 - 警告（未佩戴工服）
            'Smoking': (0, 0, 255),            # 红色 - 危险（抽烟）
            'Phone-Call': (0, 165, 255),       # 橙色 - 警告（接打电话）
            'Fire': (0, 0, 128),               # 深红 - 极度危险（火焰）
            'Smoke': (128, 128, 128),          # 灰色 - 警告（烟雾）
            'Danger-Zone': (255, 0, 255),      # 紫色 - 危险区域
            'Safety-Equipment': (0, 255, 255), # 黄色 - 安全设备
            'Vehicle': (255, 0, 0)             # 蓝色 - 车辆
        }
        
        # 风险级别定义
        self.risk_levels = {
            'Person': 'info',
            'Helmet': 'safe',
            'No-Helmet': 'high',
            'Safety-Vest': 'safe', 
            'No-Safety-Vest': 'medium',
            'Smoking': 'high',
            'Phone-Call': 'medium',
            'Fire': 'critical',
            'Smoke': 'medium',
            'Danger-Zone': 'high',
            'Safety-Equipment': 'safe',
            'Vehicle': 'info'
        }
        
        self.load_model()
    
    def load_model(self):
        """加载训练好的模型"""
        if not os.path.exists(self.model_path):
            print(f"错误: 找不到模型文件 {self.model_path}")
            print("请先训练综合安全检测模型")
            return False
        
        try:
            self.model = YOLO(self.model_path)
            print(f"成功加载综合安全检测模型: {self.model_path}")
            return True
        except Exception as e:
            print(f"加载模型失败: {e}")
            return False
    
    def detect_image(self, image_path, conf_threshold=0.5, save_result=True):
        """
        对单张图片进行综合安全检测
        
        Args:
            image_path: 图片路径
            conf_threshold: 置信度阈值
            save_result: 是否保存检测结果
        
        Returns:
            检测结果和标注图片
        """
        if self.model is None:
            print("模型未加载")
            return None
        
        # 读取图片
        image = cv2.imread(image_path)
        if image is None:
            print(f"无法读取图片: {image_path}")
            return None
        
        # 进行检测
        results = self.model(image, conf=conf_threshold)
        
        # 绘制检测结果
        annotated_image = self.draw_results(image, results[0])
        
        if save_result:
            # 保存结果
            output_dir = "comprehensive_detection_results"
            os.makedirs(output_dir, exist_ok=True)
            
            filename = Path(image_path).stem
            output_path = os.path.join(output_dir, f"{filename}_comprehensive_detected.jpg")
            cv2.imwrite(output_path, annotated_image)
            print(f"综合检测结果已保存到: {output_path}")
        
        return results[0], annotated_image
    
    def draw_results(self, image, results):
        """
        在图片上绘制检测结果（增强版本，包含风险级别指示）
        
        Args:
            image: 原始图片
            results: 检测结果
        
        Returns:
            标注后的图片
        """
        annotated_image = image.copy()
        
        if results.boxes is not None:
            for box in results.boxes:
                # 获取边界框坐标
                x1, y1, x2, y2 = box.xyxy[0].cpu().numpy().astype(int)
                
                # 获取类别和置信度
                conf = box.conf[0].cpu().numpy()
                cls = int(box.cls[0].cpu().numpy())
                class_name = self.class_names[cls]
                risk_level = self.risk_levels.get(class_name, 'info')
                
                # 选择颜色
                color = self.colors.get(class_name, (128, 128, 128))
                
                # 根据风险级别调整边框粗细
                thickness = {
                    'safe': 2,
                    'info': 2,
                    'medium': 3,
                    'high': 4,
                    'critical': 5
                }.get(risk_level, 2)
                
                # 绘制边界框
                cv2.rectangle(annotated_image, (x1, y1), (x2, y2), color, thickness)
                
                # 绘制标签（包含风险级别）
                risk_indicators = {
                    'safe': '✓',
                    'info': 'ℹ',
                    'medium': '⚠',
                    'high': '⚠',
                    'critical': '🔥'
                }
                
                indicator = risk_indicators.get(risk_level, '')
                label = f"{indicator} {class_name}: {conf:.2f}"
                label_size = cv2.getTextSize(label, cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0]
                
                # 标签背景
                cv2.rectangle(annotated_image, 
                            (x1, y1 - label_size[1] - 10), 
                            (x1 + label_size[0], y1), 
                            color, -1)
                
                # 标签文字
                cv2.putText(annotated_image, label, 
                          (x1, y1 - 5), 
                          cv2.FONT_HERSHEY_SIMPLEX, 0.6, 
                          (255, 255, 255), 2)
        
        return annotated_image
    
    def analyze_comprehensive_safety(self, results):
        """
        综合安全分析
        
        Args:
            results: 检测结果
        
        Returns:
            综合安全分析报告
        """
        if results.boxes is None:
            return {
                'total_detections': 0,
                'safety_violations': [],
                'risk_score': 0,
                'safety_status': 'No detections',
                'recommendations': []
            }
        
        detection_counts = {class_name: 0 for class_name in self.class_names}
        violations = []
        risk_score = 0
        
        for box in results.boxes:
            cls = int(box.cls[0].cpu().numpy())
            class_name = self.class_names[cls]
            conf = box.conf[0].cpu().numpy()
            risk_level = self.risk_levels.get(class_name, 'info')
            
            detection_counts[class_name] += 1
            
            # 计算风险分数
            risk_scores = {
                'safe': -5,      # 安全行为减分
                'info': 0,       # 中性
                'medium': 10,    # 中等风险
                'high': 25,      # 高风险
                'critical': 50   # 极度危险
            }
            risk_score += risk_scores.get(risk_level, 0)
            
            # 记录违规行为
            if risk_level in ['medium', 'high', 'critical']:
                violations.append({
                    'type': class_name,
                    'confidence': float(conf),
                    'risk_level': risk_level
                })
        
        # 计算整体安全状态
        if risk_score <= 0:
            safety_status = "优秀"
        elif risk_score <= 20:
            safety_status = "良好"
        elif risk_score <= 50:
            safety_status = "一般"
        elif risk_score <= 80:
            safety_status = "需要注意"
        else:
            safety_status = "危险"
        
        # 生成建议
        recommendations = self._generate_recommendations(detection_counts, violations)
        
        return {
            'total_detections': len(results.boxes),
            'detection_counts': detection_counts,
            'safety_violations': violations,
            'risk_score': max(0, risk_score),  # 确保风险分数不为负
            'safety_status': safety_status,
            'recommendations': recommendations
        }
    
    def _generate_recommendations(self, detection_counts, violations):
        """根据检测结果生成安全建议"""
        recommendations = []
        
        # 安全帽相关建议
        if detection_counts['No-Helmet'] > 0:
            recommendations.append("发现有人员未佩戴安全帽，请立即佩戴")
        
        # 工服相关建议
        if detection_counts['No-Safety-Vest'] > 0:
            recommendations.append("发现有人员未穿安全背心，建议穿戴反光背心")
        
        # 抽烟相关建议
        if detection_counts['Smoking'] > 0:
            recommendations.append("检测到吸烟行为，工地禁止吸烟，存在火灾风险")
        
        # 接打电话建议
        if detection_counts['Phone-Call'] > 0:
            recommendations.append("检测到接打电话，作业时请专注工作，注意安全")
        
        # 火焰和烟雾建议
        if detection_counts['Fire'] > 0:
            recommendations.append("⚠️ 检测到火焰！立即启动消防应急程序")
        
        if detection_counts['Smoke'] > 0:
            recommendations.append("检测到烟雾，请检查是否有火灾隐患")
        
        # 如果没有违规，给出积极反馈
        if not violations:
            recommendations.append("✅ 当前未发现明显安全违规，请继续保持")
        
        return recommendations
    
    def generate_safety_report(self, results, image_path):
        """生成详细的安全报告"""
        analysis = self.analyze_comprehensive_safety(results)
        
        report = {
            'timestamp': datetime.now().isoformat(),
            'image_path': image_path,
            'analysis': analysis,
            'report_summary': {
                'total_detections': analysis['total_detections'],
                'violations_count': len(analysis['safety_violations']),
                'risk_level': analysis['safety_status'],
                'critical_issues': [v for v in analysis['safety_violations'] if v['risk_level'] == 'critical']
            }
        }
        
        return report
    
    def save_report(self, report, output_dir="safety_reports"):
        """保存安全报告为JSON文件"""
        os.makedirs(output_dir, exist_ok=True)
        
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        filename = f"safety_report_{timestamp}.json"
        filepath = os.path.join(output_dir, filename)
        
        with open(filepath, 'w', encoding='utf-8') as f:
            json.dump(report, f, ensure_ascii=False, indent=2)
        
        print(f"安全报告已保存到: {filepath}")
        return filepath

def create_comprehensive_data_config():
    """创建综合安全检测的数据配置文件"""
    config = {
        'train': 'comprehensive_dataset/train/images',
        'val': 'comprehensive_dataset/valid/images',
        'test': 'comprehensive_dataset/test/images',
        'nc': 12,  # 类别数量
        'names': [
            'Person',
            'Helmet', 
            'No-Helmet',
            'Safety-Vest',
            'No-Safety-Vest',
            'Smoking',
            'Phone-Call',
            'Fire',
            'Smoke',
            'Danger-Zone',
            'Safety-Equipment',
            'Vehicle'
        ]
    }
    
    with open('comprehensive_data.yaml', 'w', encoding='utf-8') as f:
        import yaml
        yaml.dump(config, f, default_flow_style=False, allow_unicode=True)
    
    print("已创建综合安全检测数据配置文件: comprehensive_data.yaml")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description='综合建筑安全检测')
    parser.add_argument('--input', '-i', required=True, help='输入图片或视频路径')
    parser.add_argument('--model', '-m', default='runs/detect/comprehensive_safety/weights/best.pt', 
                       help='模型路径')
    parser.add_argument('--conf', '-c', type=float, default=0.5, help='置信度阈值')
    parser.add_argument('--save', '-s', action='store_true', help='保存检测结果')
    parser.add_argument('--report', '-r', action='store_true', help='生成安全报告')
    parser.add_argument('--create-config', action='store_true', help='创建数据配置文件')
    
    args = parser.parse_args()
    
    # 如果要求创建配置文件
    if args.create_config:
        create_comprehensive_data_config()
        return
    
    # 创建检测器
    detector = ComprehensiveSafetyDetector(args.model)
    
    input_path = args.input
    
    # 检查输入文件是否存在
    if not os.path.exists(input_path):
        print(f"错误: 找不到输入文件 {input_path}")
        return
    
    # 判断输入类型
    file_ext = Path(input_path).suffix.lower()
    image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff']
    
    if file_ext in image_extensions:
        print(f"综合安全检测图片: {input_path}")
        results, annotated_image = detector.detect_image(input_path, args.conf, args.save)
        
        if results is not None:
            # 综合安全分析
            safety_analysis = detector.analyze_comprehensive_safety(results)
            
            print("\n" + "="*60)
            print("综合安全分析报告")
            print("="*60)
            print(f"总检测数量: {safety_analysis['total_detections']}")
            print(f"安全违规数量: {len(safety_analysis['safety_violations'])}")
            print(f"风险分数: {safety_analysis['risk_score']}")
            print(f"安全状态: {safety_analysis['safety_status']}")
            
            print(f"\n检测统计:")
            for class_name, count in safety_analysis['detection_counts'].items():
                if count > 0:
                    print(f"  {class_name}: {count}")
            
            if safety_analysis['safety_violations']:
                print(f"\n安全违规详情:")
                for violation in safety_analysis['safety_violations']:
                    print(f"  - {violation['type']} (风险: {violation['risk_level']}, 置信度: {violation['confidence']:.2f})")
            
            print(f"\n安全建议:")
            for recommendation in safety_analysis['recommendations']:
                print(f"  • {recommendation}")
            
            # 生成详细报告
            if args.report:
                report = detector.generate_safety_report(results, input_path)
                detector.save_report(report)
            
            # 显示结果
            plt.figure(figsize=(15, 10))
            plt.imshow(cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB))
            plt.title('综合建筑安全检测结果')
            plt.axis('off')
            plt.show()
    
    else:
        print(f"当前版本暂不支持视频检测，支持的图片格式: {image_extensions}")

if __name__ == "__main__":
    main() 